from flask import Flask, request, jsonify

import llm
import my_bert
import rf, my_fasttext
# 将my_bert作为主程序运行时保存的model，joblib将model注册为__main__.MyModel，而不是my_bert.MyModel，这里需要显式导入才能正常加载model
from my_bert import MyModel

app = Flask(__name__)


@app.route('/predict', methods=['post'])
def predict():
    json_data = request.get_json()
    print(f'request param:{json_data}')
    model = json_data['model']
    text = json_data['text']
    if model == 'rf':
        result = rf.predict(text)
    elif model == 'fasttext':
        result = my_fasttext.predict(text)
    elif model == 'bert':
        result = my_bert.predict(text)
    elif model == 'llm':
        result = llm.predict(text)
    else:
        raise ValueError('model type error')
    return jsonify({'result': result})


if __name__ == '__main__':
    app.run(port=80)
